#!/usr/bin/env python
import sys, os, os.path

import threading
import logging

from py_interface import erl_node, erl_eventhandler
from py_interface.erl_common import AlignShortNodeName, DebugOnAll
from py_interface.erl_opts import ErlNodeOpts
from py_interface.erl_term import *

from amqplib.client_0_8 import channel as amqp

class RabbitMQ(object):
    logger = logging.getLogger('rabbitmq')    
    
    def __init__(self, name, cookie=None, nodeName=None):
        if cookie is None:
            cookie = self.readCookie()
            
        self.node = erl_node.ErlNode(name, ErlNodeOpts(cookie=cookie))
        
        self.nodeName = nodeName or AlignShortNodeName('rabbit')
        
        self.logger.info("create node <%s> with cookie '%s'", name, cookie)
        
        self.mbox = self.node.CreateMBox(None)
        self.mbox.RegisterName('p')
        
        self.logger.info("create mbox and register as %s", 'p')
        
    def readCookie(self):
        if sys.platform == 'win32':
            filename = os.path.join(os.environ['USERPROFILE'], '.erlang.cookie')
        else:
            filename = '~/.erlang.cookie'
            
        with open(os.path.join(filename), 'r') as f:
            cookie = f.read().strip()
            
            self.logger.info("automatic read cookie '%s' from %s", cookie, filename)
            
            return cookie
            
    @property
    def alive(self):
        result = []
        
        def handler(value):
            result.append(value)
            
            erl_eventhandler.GetEventHandler().StopLooping()
            
        self.node.Ping(self.nodeName, handler)
        
        erl_eventhandler.GetEventHandler().Loop()
        
        return result == ['pong']
            
    def call(self, mod, func, args=[], callback=None):
        results = []
        
        def handler(msg):
            if callback:
                callback(msg)
            
            if type(msg) == list:
                results.extend(msg)
            else:
                results.append(msg)
            
            erl_eventhandler.GetEventHandler().StopLooping()
        
        self.mbox.SendRPC(self.nodeName, mod, func, args, handler)
        
        erl_eventhandler.GetEventHandler().Loop()
        
        return results    
        
    def stop(self):
        """Stops the Erlang node on which RabbitMQ is running."""
        self.logger.info("stop server")
        
        result = self.call('rabbit', 'stop_and_halt')
        
        return result and IsErlAtom(result[0]) and result[0].atomText == 'ok'
        
    def start_app(self):
        """Starts the RabbitMQ application."""
        self.logger.info("start application")
        
        result = self.call('rabbit', 'start')
        
        return result and IsErlAtom(result[0]) and result[0].atomText == 'ok'
    
    def stop_app(self):
        """Stops the RabbitMQ application, leaving the Erlang node running."""
        self.logger.info("stop application")

        result = self.call('rabbit', 'stop')
        
        return result and IsErlAtom(result[0]) and result[0].atomText == 'ok'
    
    def reset(self, force=False):
        """Return a RabbitMQ node to its virgin state.
        
        Removes the node from any cluster it belongs to, removes all data
        from the management database, such as configured users and vhosts,
        and deletes all persistent messages.

        The force_reset command differs from reset in that it resets
        the node unconditionally, regardless of the current management
        database state and cluster configuration. It should only be used as
        a last resort if the database or cluster configuration has been corrupted.

        For reset and force_reset to succeed the RabbitMQ application must
        have been stopped, e.g. with stop_app.
        """
        self.logger.info("reset database")
        
        result = self.call('rabbit_mnesia', 'force_reset' if force else 'reset')
        
        return result and IsErlAtom(result[0]) and result[0].atomText == 'ok'
        
    def status(self):
        """Get various information about the RabbitMQ broker,
        such as whether the RabbitMQ application on the current node,
        its version number, what nodes are part of the broker,
        which of these are running.
        """
        self.logger.info("query node status")
        
        status = { 'apps' : {}, 'nodes' : {} }
                        
        for key, values in self.call('rabbit', 'status'):
            if key.atomText == 'running_applications':
                for app, name, ver in values:
                    status['apps'][app.atomText] = (name, ver)
            elif key.atomText == 'nodes':
                for node in values:
                    status['nodes'][node.atomText] = False
            elif key.atomText == 'running_nodes':
                for node in values:
                    status['nodes'][node.atomText] = True
        
        return status
    
    def check_login(self, username, password):            
        self.logger.info("login as user %s", username)
                    
        result = self.call('rabbit_access_control', 'check_login',
                           [ErlBinary('PLAIN'), ErlBinary(username + '\0' + password)])
        
        return result and type(result[0]) == tuple and \
               result[0][0].atomText == 'user' and result[0][1].contents == username
    
    def add_user(self, username, password):
        """
        username
            The name of the user to create.
        password
            The password the created user will use to log in to the broker.
        """
        self.logger.info("add user %s", username)
        
        result = self.call('rabbit_access_control', 'add_user',
                           [ErlBinary(username), ErlBinary(password)])
        
        return result and IsErlAtom(result[0]) and result[0].atomText == 'ok'
        
    def delete_user(self, username):
        """
        username
            The name of the user to delete.
        """
        self.logger.info("delete user %s", username)
            
        result = self.call('rabbit_access_control', 'delete_user',
                           [ErlBinary(username)])
        
        return result and IsErlAtom(result[0]) and result[0].atomText == 'ok'
        
    def change_password(self, username, password):
        """
        username
            The name of the user whose password is to be changed.
        newpassword
            The new password for the user.
        """
        self.logger.info("change user %s password", username)
        
        result = self.call('rabbit_access_control', 'change_password',
                           [ErlBinary(username), ErlBinary(password)])
        
        return result and IsErlAtom(result[0]) and result[0].atomText == 'ok'
    
    def list_users(self):
        """list all users."""
        
        self.logger.info("list users")
            
        return [user.contents if IsErlBinary(user) else str(user)
                for user in self.call('rabbit_access_control', 'list_users')]
        
    def add_vhost(self, vhostpath):
        """
        vhostpath
            The name of the virtual host entry to create.
        """
        self.logger.info("add vhost <%s>", vhostpath)
        
        result = self.call('rabbit_access_control', 'add_vhost', [ErlBinary(vhostpath)])
        
        return result and IsErlAtom(result[0]) and result[0].atomText == 'ok'
    
    def delete_vhost(self, vhostpath):
        """
        vhostpath
            The name of the virtual host entry to delete.
            
        Deleting a virtual host deletes all its exchanges, queues, user mappings and associated permissions.
        """
        self.logger.info("delete vhost <%s>", vhostpath)
        
        result = self.call('rabbit_access_control', 'delete_vhost', [ErlBinary(vhostpath)])
        
        return result and IsErlAtom(result[0]) and result[0].atomText == 'ok'
    
    def list_vhosts(self):
        self.logger.info("list virtual hosts")
            
        return [vhost.contents if IsErlBinary(vhost) else str(vhost)
                for vhost in self.call('rabbit_access_control', 'list_vhosts')]
        
    def to_python(self, value):
        if IsErlAtom(value):
            return { 'true' : True, 'false' : False }.get(value.atomText, value.atomText)
        if IsErlBinary(value):
            return value.contents
        if type(value) == list:
            return list(self.to_python(v) for v in value)
        if type(value) == tuple:
            return tuple(self.to_python(v) for v in value)
        
        return value
        
    def list_queues(self, vhostpath='/'):
        self.logger.info("list queues of virtual host <%s>", vhostpath)
        
        result = {}
        
        for entry in self.call('rabbit_amqqueue', 'info_all', [ErlBinary(vhostpath)]):
            queue = {}
            
            for name, value in entry:
                if name.atomText == 'name':
                    type, vhostpath, kind, name = self.to_python(value)
                    
                    queue['name'] = name
                else:
                    queue[name.atomText] = self.to_python(value)
                
            result[queue['name']] = queue
        
        return result

import unittest

mq = None

class TestRabbitMQ(unittest.TestCase):        
    def testInit(self):
        self.assert_(mq.alive)
        
    def testStatus(self):
        status = mq.status()
        
        self.assert_(status['nodes'][AlignShortNodeName('rabbit')])
        self.assert_(status['apps'].has_key('rabbit'))
        self.assert_(status['apps'].has_key('mnesia'))            
            
    def testUsers(self):
        users = mq.list_users()
        
        self.assert_('guest' in users)
        
        self.assertFalse('test' in users)
        
        self.assert_(mq.add_user('test', 'test'))
        
        self.assert_('test' in mq.list_users())
        
        self.assert_(mq.check_login('test', 'test'))
                
        self.assert_(mq.delete_user('test'))
        
        self.assertFalse('test' in users)
        
    def testVHosts(self):
        self.assert_('/' in mq.list_vhosts())
        
        self.assert_(mq.add_vhost('/test'))
        
        self.assert_('/test' in mq.list_vhosts())
        
        self.assert_(mq.delete_vhost('/test'))
        
        self.assertFalse('/test' in mq.list_vhosts())
        
    def testQueue(self):
        self.assertFalse('test' in mq.list_queues().keys())

if __name__ == '__main__':
    debugOn = '-v' in sys.argv
    
    logging.basicConfig(level=logging.DEBUG if debugOn else logging.WARN,
                        format='%(asctime)s %(levelname)s %(name)s %(message)s')
    
    if debugOn:
        DebugOnAll()
        
    mq = RabbitMQ('test')
    
    unittest.main()